#version 450

#extension GL_GOOGLE_include_directive : require
#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require

#include "../Brdf.glsl"
#include "../Lighting.glsl"


layout (location = 0) in vec2 inUV;


layout (location = 0) out vec4 outColor;


const float CosineA0 = M_PI;
const float CosineA1 = (2.0 * M_PI) / 3.0;
const float CosineA2 = M_PI * 0.25;
const float IndirectSpecularStrength = 2.0f;


layout (set = 0, binding = 0) uniform sampler2D gBuffer1; // RGB: Albedo, A: Metallic
layout (set = 0, binding = 1) uniform sampler2D gBuffer2; // RG: Normal, BA: Motion Vector
layout (set = 0, binding = 2) uniform sampler2D gBuffer3; // R: Roughness, G: Curvature, B: Mesh ID, A: Linear Z
layout (set = 0, binding = 3) uniform sampler2D gBufferDepth;

layout (set = 1, binding = 0) uniform sampler2D shadowImage;

layout (set = 2, binding = 0) uniform sampler2D aoImage;

layout (set = 3, binding = 0) uniform sampler2D giImage;

layout (set = 4, binding = 0) uniform sampler2D reflectionsImage;

layout (set = 5, binding = 0) uniform PerFrameUBO
{
    mat4 viewInverse;
    mat4 projInverse;
    mat4 projViewInverse;
    mat4 prevProjView;
    mat4 projView;
    vec4 cameraPosition;
    vec4 currentPrevJitter;
    Light light;
} perFrameUBO;

layout (set = 5, binding = 1) uniform sampler2D shImage;
layout (set = 5, binding = 2) uniform samplerCube prefilterImage;
layout (set = 5, binding = 3) uniform sampler2D brdfLutImage;


layout (push_constant) uniform PushConstants
{
    int shadow;
    int ao;
    int gi;
    int reflections;
} pushConstants;


struct SH9
{
    float c[9];
};


struct SH9Color
{
    vec3 c[9];
};


void ProjectOntoSh9(in vec3 dir, inout SH9 sh)
{
    // Band 0
    sh.c[0] = 0.282095;

    // Band 1
    sh.c[1] = -0.488603 * dir.y;
    sh.c[2] = 0.488603 * dir.z;
    sh.c[3] = -0.488603 * dir.x;

    // Band 2
    sh.c[4] = 1.092548 * dir.x * dir.y;
    sh.c[5] = -1.092548 * dir.y * dir.z;
    sh.c[6] = 0.315392 * (3.0 * dir.z * dir.z - 1.0);
    sh.c[7] = -1.092548 * dir.x * dir.z;
    sh.c[8] = 0.546274 * (dir.x * dir.x - dir.y * dir.y);
}


vec3 EvaluateSh9Irradiance(in vec3 direction)
{
    SH9 basis;

    ProjectOntoSh9(direction, basis);

    basis.c[0] *= CosineA0;
    basis.c[1] *= CosineA1;
    basis.c[2] *= CosineA1;
    basis.c[3] *= CosineA1;
    basis.c[4] *= CosineA2;
    basis.c[5] *= CosineA2;
    basis.c[6] *= CosineA2;
    basis.c[7] *= CosineA2;
    basis.c[8] *= CosineA2;

    vec3 color = vec3(0.0);

    for (int i = 0; i < 9; i++)
        color += texelFetch(shImage, ivec2(i, 0), 0).rgb * basis.c[i];

    color.x = max(0.0, color.x);
    color.y = max(0.0, color.y);
    color.z = max(0.0, color.z);

    return color / M_PI;
}


vec3 FresnelSchlickRoughness(float cosTheta, vec3 F0, float roughness)
{
    return F0 + (max(vec3(1.0 - roughness), F0) - F0) * pow(max(1.0 - cosTheta, 0.0), 5.0);
}


vec3 IndirectLighting(vec3 N, vec3 diffuseColor, float roughness, float metallic, float ao, vec3 Wo, vec3 F0)
{
    const vec3 R = reflect(-Wo, N);

    vec3 F = FresnelSchlickRoughness(max(dot(N, Wo), 0.0), F0, roughness);

    vec3 kS = F;
    vec3 kD = 1.0 - kS;
    kD *= 1.0 - metallic;

    vec3 irradiance = pushConstants.gi == 1 ? textureLod(giImage, inUV, 0.0f).rgb : EvaluateSh9Irradiance(N);
    vec3 diffuse = irradiance * diffuseColor;

    const float MAX_REFLECTION_LOD = 4.0;
    vec3 prefilteredColor = pushConstants.reflections == 1 ? textureLod(reflectionsImage, inUV, 0.0f).rgb : textureLod(prefilterImage, R, roughness * MAX_REFLECTION_LOD).rgb;
    vec2 brdf = texture(brdfLutImage, vec2(max(dot(N, Wo), 0.0), roughness)).rg;
    vec3 specular = prefilteredColor * (F * brdf.x + brdf.y) * IndirectSpecularStrength;

    return (kD * diffuse + specular) * ao;
}


void main()
{
    vec4 gBufferData1 = texture(gBuffer1, inUV);
    vec4 gBufferData2 = texture(gBuffer2, inUV);
    vec4 gBufferData3 = texture(gBuffer3, inUV);

    

    const vec3 worldPos = WorldPositionFromDepth(inUV, texture(gBufferDepth, inUV).r, perFrameUBO.projViewInverse);
    const vec3 albedo = gBufferData1.rgb;
    const float metallic = gBufferData1.a;
    const float roughness = gBufferData3.r;
    const float visibility = pushConstants.shadow == 1 ? texture(shadowImage, inUV).r : 1.0f;
    const float ao = pushConstants.ao == 1 ? texture(aoImage, inUV).r : 1.0f;


//    outColor = vec4(visibility, visibility, visibility, 1.0);
//    return;


    const vec3 N = OctohedralToDirection(gBufferData2.rg);
    const vec3 Wo = normalize(perFrameUBO.cameraPosition.xyz - worldPos);

    const vec3 F0 = mix(vec3(0.04f), albedo, metallic);
    const vec3 diffuseColor = mix(albedo * (vec3(1.0f) - F0), vec3(0.0f), metallic);

    vec3 Lo = vec3(0.0f);

    // Direct Lighting
    Lo += DirectLighting(perFrameUBO.light, Wo, N, worldPos, F0, diffuseColor, roughness) * visibility;

    // Indirect lighting
    Lo += IndirectLighting(N, diffuseColor, roughness, metallic, ao, Wo, F0);

    outColor = vec4(Lo, 1.0);
}